import json
import os
import uuid
import argparse

from models.experimental import attempt_load

from PIL import Image
from flask import Flask, request, jsonify
from flask_cors import CORS, cross_origin
from flask_sqlalchemy import SQLAlchemy
import base64

from my import delExp
from utils.torch_utils import select_device
from service.RecordService import RecordRepository
from model.Record import Record
from detect import predict
from flask_sslify import SSLify

parser = argparse.ArgumentParser()
parser.add_argument('--weights', nargs='+', type=str, default='best.pt', help='model.pt path(s)')
parser.add_argument('--no-trace', action='store_true', help='don`t trace model')
parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
parser.add_argument('--img-size', type=int, default=800, help='inference size (pixels)')

app = Flask(__name__)
sslify = SSLify(app)  # 用于强制将 HTTP 请求重定向到 HTTPS。
opt = parser.parse_args()  # 解析命令行参数，将结果存储在 opt 变量中。
device = select_device(opt.device)
trace = not opt.no_trace  # 取反 是否启用追踪模型。
CORS(app)  # 配置了 CORS（跨域资源共享）支持
UPLOAD_FOLDER = 'uploads'
app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
weights = opt.weights
model = attempt_load(weights, map_location=device)
model.half()  # 半精度浮点数格式
# 确保上传文件夹存在
if not os.path.exists(app.config['UPLOAD_FOLDER']):
    os.makedirs(app.config['UPLOAD_FOLDER'])

# 配置数据库连接
app.config['SQLALCHEMY_DATABASE_URI'] = 'mysql+pymysql://root:1122334.1@127.0.0.1/helmet?charset=utf8mb4'
app.config['SQLALCHEMY_TRACK_MODIFICATIONS'] = False  # 指定是否追踪对象的修改并发送信号

# 初始化SQLAlchemy
db = SQLAlchemy(app)
record_repository = RecordRepository(db.session)


# 读取本地文件并转换为Base64编码
def file_to_base64(file_path):
    with open(file_path, "rb") as file:
        encoded_string = base64.b64encode(file.read())
        return encoded_string


class User(db.Model):
    id = db.Column(db.Integer, primary_key=True)
    username = db.Column(db.String(255), unique=True, nullable=False)
    password = db.Column(db.String(255), unique=False, nullable=False)
    bz = db.Column(db.String(255), unique=False, nullable=False)

    def __repr__(self):
        return f'<User {self.username}>'

    def check_password(self, password):
        return self.password == password

    def get_all(self):
        return self.session.query(Record).order_by(Record.curtime.desc()).all()

    def to_dict(self):
        # 注意：通常不应该返回密码哈希或敏感信息
        return {
            'id': self.id,
            'username': self.username,
            'bz': self.bz,
            'password':self.password
        }


with app.app_context():
    db.create_all()



@app.route('/upload2', methods=['POST'])    # 处理视频
def upload_file2():
    data = request.json.get('file').split(',')[1]  # 假设base64编码数据作为JSON请求的一个字段传递过来
    print(data)
    if data:
        file_path = os.path.join('./inference/images/', 'cur.jpg')
        # 解码base64数据并写入文件
        with open(file_path, 'wb') as file:
            file.write(base64.b64decode(data))
    img = Image.open('inference/images/cur.jpg')
    # 强制调整图片尺寸为500x333像素
    img = img.resize((500, 333), Image.ANTIALIAS)
    if img.mode == 'RGBA':
        img = img.convert('RGB')
    img.save('inference/images/cur.jpg')
    record_list = predict(model)
    delExp()
    records_json = [record.to_dict() for record in record_list]
    return jsonify({'image_url': '/uploads/cur.jpg', 'draw_url': str(file_to_base64('./upload/res.jpg')),
                    'list': json.dumps(records_json)})


@app.route('/upload', methods=['POST'])     # 处理图片
def upload_file():
    # 检查是否有文件在请求中
    if 'file' not in request.files:
        return jsonify({'error': 'No file part in the request'}), 400

    file = request.files['file']

    # 如果用户没有选择文件，浏览器也会提交一个空文件部分，没有文件名
    if file.filename == '':
        return jsonify({'error': 'No selected file'}), 400

        # 读取文件内容

    # 生成一个唯一的文件名
    filename = str(uuid.uuid4()) + '.' + file.filename.split('.')[
        -1]  # 根据UUID生成一个唯一的文件名，通过拆分文件名获取文件扩展名，并将唯一文件名和扩展名拼接起来，生成最终的文件名
    file_path = os.path.join(app.config['UPLOAD_FOLDER'], filename)

    # 保存文件到本地
    file.save('inference/images/cur.jpg')
    # 打开上传的图片文件
    img = Image.open('inference/images/cur.jpg')
    # 强制调整图片尺寸为500x333像素
    img = img.resize((500, 333), Image.ANTIALIAS)
    if img.mode == 'RGBA':
        img = img.convert('RGB')
    img.save('inference/images/cur.jpg')

    # if trace:
    # model = TracedModel(model, device, opt.img_size)
    record_list = predict(model)
    delExp()
    records_json = [record.to_dict() for record in record_list]
    return jsonify({'image_url': '/uploads/cur.jpg', 'draw_url': str(file_to_base64('./upload/res.jpg')),
                    'list': json.dumps(records_json)})


# @app.route('/uploads/<filename>')
# def uploaded_file(filename):
#     # 检查文件名是否安全，防止路径遍历攻击
#     if '..' in filename or filename.startswith('/'):
#         abort(404)
#     return send_from_directory(app.config['UPLOAD_FOLDER'], filename)


@app.route('/login', methods=['POST'])
def login():
    data = request.get_json()
    username = data.get('username')
    password = data.get('password')

    if not username or not password:
        return jsonify({'error': 'Username and password are required'}), 400

    user = User.query.filter_by(username=username).first()
    if user and user.check_password(password):
        # 如果用户名和密码都正确，返回用户信息
        return jsonify(user.to_dict()), 200
    else:
        # 如果用户名或密码不正确，返回错误
        return jsonify({'error': 'Invalid username or password'}), 401


@app.route('/register', methods=['POST'])
def register():
    data = request.get_json()
    username = data.get('username')
    password = data.get('password')
    # 创建新的User对象
    new_user = User(username=username, password=password, bz='')

    # 将新对象添加到数据库会话中
    db.session.add(new_user)

    # 提交会话以保存更改到数据库
    try:
        db.session.commit()
        print(f"User {username} has been added to the database.")
        return jsonify(f"User {username} has been added to the database."), 200
    except Exception as e:
        # 如果发生错误，回滚会话
        db.session.rollback()
        print(f"An error occurred: {e}")
        print("User was not added to the database.")
        return jsonify({'error': '注册失败'}), 401


@app.route('/users/count', methods=['GET'])
def get_users_count():
    users_count = User.query.count()
    return jsonify({'count': users_count})


@app.route('/userlist', methods=['GET'])
def get_users_list():
    userList = User.query.all()
    records_json = [record.to_dict() for record in userList]
    return json.dumps(records_json)


@app.route('/users/<int:user_id>', methods=['PUT'])
def update_user(user_id):
    user = User.query.get(user_id)
    if not user:
        return jsonify({'message': 'User not found'}), 404

    data = request.get_json()
    user.username = data.get('username', user.username)
    user.password = data.get('password', user.password)
    user.bz = data.get('bz', user.bz)
    db.session.commit()

    return jsonify({'message': 'User updated successfully'})


@app.route('/users/<int:user_id>', methods=['DELETE'])
def delete_user(user_id):
    user = User.query.get(user_id)
    if not user:
        return jsonify({'message': 'User not found'}), 404

    db.session.delete(user)
    db.session.commit()

    return jsonify({'message': 'User deleted successfully'})


@cross_origin()
@app.route('/addRecord', methods=['POST'])
def add_record():
    data = request.get_json()
    new_record = Record(label=data['label'], position=data['position'], zxd=data['zxd'], user=data['user'],
                        curtime=None, img=None, bz=data['bz'])
    record_repository.add(new_record)
    return jsonify('添加记录成功'), 200


@app.route('/delRecord', methods=['POST'])
def delRecord():
    data = request.get_json()
    record_repository.delete(data['id'])
    return jsonify('删除记录成功'), 200


@app.route('/update', methods=['POST'])
def update():
    data = request.get_json()
    new_record = Record(id=data['id'], label=data['label'], position=data['position'], zxd=data['zxd'],
                        user=data['user'],
                        curtime=None, img=base64.b64decode(data['img']), bz=data['bz'])
    record_repository.update(new_record)
    return jsonify('更新记录成功'), 200


@app.route('/recordList', methods=['GET'])
def recordList():
    records = record_repository.get_all()
    records_json = [record.to_dict() for record in records]
    return json.dumps(records_json)


if __name__ == '__main__':
    app.run(debug=True, port=5003, host='0.0.0.0')
